//
// Created by Lenovo on 2025/8/22.
//

#include "PLTHook.h"
#include <linux/elf.h>
#include "Log.h"
#include <sys/mman.h>
#include <sys/socket.h>


void PLTHook::setNewFuncPtr(uintptr_t newFuncPtr) {
    this->newFuncPtr = newFuncPtr;
}

uintptr_t PLTHook::getOriginFuncPtr() {
    return this->originFuncPtr;
}


intptr_t PLTHook::parseSOBaseAddress(const std::string soName) {
    FILE *file = fopen("/proc/self/maps", "r");
    char line[2014];
    uintptr_t baseAddr = 0;
    while (fgets(line, sizeof(line), file)) {
        if (nullptr != strstr(line, soName.c_str())) {  // libcppdemo.so
            std::string targetLine = line;
            std::size_t splitPosition = targetLine.find('-');
            if (splitPosition != std::string::npos) {
                std::string startAddr = targetLine.substr(0, splitPosition);
                baseAddr = std::strtoull(startAddr.c_str(), nullptr, 16);
                break;
            }
        }
    }
    return baseAddr;
}

/**
 * @param soBaseAddr SO文件加载至虚拟内存中的基地址
 * @return 程序头表中.dynamic segment的起始地址及其所占的大小
 */
std::pair<uintptr_t, uint64_t> PLTHook::getDynAddress(uintptr_t soBaseAddress) {
    Log::debug(TAG, "so base address is %x", soBaseAddress);
    auto *ehdr = (Elf64_Ehdr *) soBaseAddress;
    uint16_t programHeaderCount = ehdr->e_phnum;
    uint64_t programHeaderOffset = ehdr->e_phoff;
    Log::debug(TAG, "number of program header size is %x", programHeaderCount);
    Log::debug(TAG, "offset of program header size is %llx", programHeaderOffset);
    auto *pProgramHeader = reinterpret_cast<Elf64_Phdr *>(soBaseAddress + programHeaderOffset);
    for (uint16_t index = 0; index < programHeaderCount; index++) {
        uint16_t segmentType = pProgramHeader[index].p_type;
        if (segmentType == PT_DYNAMIC) {
            Log::debug(TAG, "file offset of dynamic is %llx", pProgramHeader[index].p_offset);
            Log::debug(TAG, "virtual address of dynamic is %llx", pProgramHeader[index].p_vaddr);
            Log::debug(TAG, "size of dynamic is %llx", pProgramHeader[index].p_memsz);
            return std::make_pair(soBaseAddress + pProgramHeader[index].p_vaddr,
                                  pProgramHeader[index].p_memsz);
        }
    }
    return std::make_pair(0, 0);
}

/**
 * @param pStrTbl 动态链接字符串表的起始地址
 * @param strTblSize 动态链接字符串表所占的大小
 */
void PLTHook::printDynSymName(uintptr_t pStrTbl, uint64_t strTblSize) {
    uint64_t index = 0;
    uint64_t symNameCount = 0;
    while (index < strTblSize) {
        char *symName = reinterpret_cast<char *>(pStrTbl) + index;
        size_t len = strlen(symName) + 1;
        Log::debug(TAG, "index = %lld, symName = %s, len= %zu", index, symName, len);
        index = index + len;
        symNameCount++;
    }
    Log::debug(TAG, "symNameCount = %llu", symNameCount);
}

/**
 * 获取PLT所依赖的Rela，该结构中包含指向GOT表中目标函数的地址
 * @param soBaseAddress
 * @param dynAddress
 * @return
 */
uintptr_t PLTHook::parseRELPLT(uintptr_t soBaseAddress, std::pair<uintptr_t, uint64_t> dynAddress) {
    auto *dyn = (Elf64_Dyn *) dynAddress.first;
    uint64_t memSize = dynAddress.second;
    uint64_t num = memSize / sizeof(Elf64_Dyn);  // equals readelf.exe -d .\libcppdemo.so
    uintptr_t pRelPltTbl = 0;
    uint64_t relPltTblSize = 0;
    uint64_t relPltTblItemType = 0;
    uintptr_t pRelaTbl = 0;
    uint64_t relaSize = 0;
    uintptr_t pDynSymTbl = 0;
    uintptr_t pStrTbl = 0;
    uint64_t strTblSize = 0;
    Log::debug(TAG, "number of dynamic is %lld", num);
    for (auto index = 0; index < num; index++) {
        Log::debug(TAG, "type of dynamic is %lld", dyn[index].d_tag);
        switch (dyn[index].d_tag) {
            case DT_JMPREL: // rela.plt
                pRelPltTbl = soBaseAddress + dyn[index].d_un.d_ptr;
                Log::debug(TAG, "pRelPltTbl=%u", pRelPltTbl);
                break;
            case DT_PLTRELSZ:
                relPltTblSize = dyn[index].d_un.d_val;
                Log::debug(TAG, "relPltTblSize=%llu", relPltTblSize);
                break;
            case DT_PLTREL:
                relPltTblItemType = dyn[index].d_un.d_val;
                Log::debug(TAG, "dyn[index].d_un.d_val = %llu", relPltTblItemType);
                break;
            case DT_RELA:
                pRelaTbl = soBaseAddress + dyn[index].d_un.d_ptr;
                Log::debug(TAG, "pRelaTbl=%u", pRelaTbl);
                break;
            case DT_RELASZ:
                relaSize = dyn[index].d_un.d_val;
                Log::debug(TAG, "relaSize=%llu", relaSize);
                break;
            case DT_SYMTAB:
                pDynSymTbl = soBaseAddress + dyn[index].d_un.d_ptr;
                break;
            case DT_STRTAB:
                pStrTbl = soBaseAddress + dyn[index].d_un.d_ptr;
                break;
            case DT_STRSZ:
                strTblSize = dyn[index].d_un.d_val;
                break;
            default:
                break;
        }
    }
    printDynSymName(pStrTbl, strTblSize);

    uint64_t relPltTblItemCount =
            relPltTblSize / (relPltTblItemType == DT_RELA ? sizeof(Elf64_Rela) : sizeof(Elf64_Rel));
    Log::debug(TAG, "count of relPltTblItemCount is %lld", relPltTblItemCount);
    if (pRelPltTbl == 0 || pDynSymTbl == 0 || pStrTbl == 0) {
        Log::error(TAG, "pRelPltTbl、pDynSymTbl or pStrTbl is 0");
        return 0;
    }
    for (int index = 0; index < relPltTblItemCount; index++) {
        Log::debug(TAG, "(Elf64_Rela *) pRelPltTbl is %u", (Elf64_Rela *) pRelPltTbl);
        uint64_t dynSymInfo = (relPltTblItemType == DT_RELA ?
                               (reinterpret_cast<Elf64_Rela *>(pRelPltTbl))[index].r_info
                                                            : (reinterpret_cast<Elf64_Rel *>(pRelPltTbl))[index].r_info);
        uint64_t dynSymIdx = ELF64_R_SYM(dynSymInfo);
        Log::debug(TAG, "index of DynSymTbl is %lld", dynSymIdx);
        uint64_t synSymNameIdx = ((reinterpret_cast<Elf64_Sym *>(pDynSymTbl))[dynSymIdx]).st_name;
        char *synSymName = (reinterpret_cast<char *>(pStrTbl)) + synSymNameIdx;
        if (strcmp("malloc", synSymName) == 0) {
            Log::debug(TAG, "synSymName is %s", synSymName);
            uintptr_t gotEntryAddress = soBaseAddress + (relPltTblItemType == DT_RELA ?
                                             ((Elf64_Rela *) pRelPltTbl)[index].r_offset +
                                             ((Elf64_Rela *) pRelPltTbl)[index].r_addend
                                                                          : ((Elf64_Rel *) pRelPltTbl)[index].r_offset);

            uintptr_t pageStart = gotEntryAddress & ~(PAGE_SIZE - 1);
            if (mprotect(reinterpret_cast<void *>(pageStart), PAGE_SIZE, PROT_READ | PROT_WRITE) <
                0) {
                Log::debug(TAG, "mprotect failed, %s", strerror(errno));
                return 0;
            }
            // 保存原函数指针
            originFuncPtr = reinterpret_cast<uintptr_t>(*reinterpret_cast<void **>(gotEntryAddress));
            // 替换原函数指针
            *reinterpret_cast<void **>(gotEntryAddress) = reinterpret_cast<void *>(newFuncPtr);
            __builtin___clear_cache(reinterpret_cast<char *>(pageStart),
                                    reinterpret_cast<char *>(pageStart + sizeof(void *)));
            return originFuncPtr;
        }
    }
    return 0;
}
